

clearvars -except diffRadius


%% data and prior
load('../../true_DGP/priorWSigma.mat')
data=readtable('../../../Data/data_main_sample.csv');
dd=readtable('../../../Data/distance_capitals.csv');

data=data(data.year>=1950&data.year<=2000&isfinite(data.D),:);
ccodes=unique(data.ccode);
years=unique(data.year(data.year>1950));
N=length(ccodes);
T=length(years);
y=zeros(N,T); D=y; M=y;
for n=1:N
    dn=data(data.ccode==ccodes(n),:);
    for t=1:T
        if isempty(setdiff(years(t),dn.year))
            if isfinite(dn.y(years(t)==dn.year))
            	M(n,t)=1; % not missing
                y(n,t)=dn.y(years(t)==dn.year);
                D(n,t)=dn.D(years(t)==dn.year);
            end
        end
    end
end

Z=10^(-6)*table2array(dd(:,2:end));
KZ=size(Z,3);

% democracy diffusion/contagion
X=zeros(N,T);
nm0=1*arrayfun(@(n)isempty(setdiff(ccodes(n),data.ccode(data.year==1950&isfinite(data.D)))),1:length(ccodes))';
D0=arrayfun(@(n)data.D(data.year==1950&data.ccode==ccodes(n)),1:length(ccodes),'UniformOutput',0)';
D0(nm0==0)={0};
D0=cell2mat(D0);
for n=1:N
    for t=1:T
        if M(n,t)==1
            if t==1
                weights=nm0.*exp(-diffRadius*Z(:,n));
                weights(n)=0;
                weights=weights/sum(weights);
                X(n,t)=weights'*D0;
            else
                weights=M(:,t-1).*exp(-diffRadius*Z(:,n));
                weights(n)=0;
                weights=weights/sum(weights);
                X(n,t)=weights'*D(:,t-1);
            end
        end
    end
end
X=M.*(X-sum(sum(X))/sum(sum(M)));
K=1;
clear nm0 D0 t weights

ep_sd=sqrt(diag(Sigma));

% Sparse-grid integration
[sgi_nodes,sgi_weights]=nwspgr('KPN',2,8);

cnames=cell(N,1);
for n=1:N
    a=unique(data.idacr(data.ccode==ccodes(n)));
    cnames{n}=a{1};
end

ccodes=table(ccodes,cnames);

clear a cnames data dd n


%% for Knitro

% Jacobian sparsity pattern:
JPattern=[repmat(eye(N),T,1),ones(N*T,2),ones(N*T,2*N),ones(N*T,N),...
    repmat(eye(N),T,1),zeros(N*T,N),ones(N*T,K),ones(N*T,KZ),eye(N*T)];
JPattern=sparse(JPattern);

% Hessian sparsity pattern
% log-posterior
HPattern=zeros(6*N+2+K+KZ+N*T,6*N+2+K+KZ+N*T);
  % alpha,theta,beta,v,f
HPattern(1:(5*N+2),1:(5*N+2))=eye(5*N+2);
  % vs
HPattern(5*N+2+(1:N),:)=[zeros(N,5*N+2),eye(N),zeros(N,K+KZ),repmat(eye(N),1,T)];
  % kappa
HPattern((6*N+2+K+KZ+1):end,:)=[zeros(N*T,5*N+2),repmat(eye(N),T,1),zeros(N*T,K+KZ),eye(N*T)];
% equilibrium constraints
  % alpha
  HPattern(1:N,:)=HPattern(1:N,:)+...
        [repmat(eye(N),1,2),ones(N,2*N+2),eye(N),zeros(N),ones(N,K+KZ),repmat(eye(N),1,T)];
  % theta
  HPattern(N+(1:2),:)=HPattern(N+1,:)+ones(2,6*N+2+K+KZ+N*T);
  % f
  HPattern(4*N+(1:N),:)=HPattern(4*N+(1:N),:)+...
        [repmat(eye(N),1,2),ones(N,2*N+2),eye(N),zeros(N),ones(N,K+KZ),repmat(eye(N),1,T)];
  % beta,v,xi,gamma
for n=2:4
    HPattern((n-1)*N+2+(1:N),:)=HPattern((n-1)*N+2+(1:N),:)+...
        [ones(N,5*N+2),zeros(N),ones(N,K+KZ+N*T)];
end
HPattern(6*N+2+(1:K),:)=HPattern(6*N+2+(1:K),:)+...
        [ones(K,5*N+2),zeros(K,N),ones(K,K+KZ+N*T)];
HPattern(6*N+2+K+(1:KZ),:)=HPattern(6*N+2+K+(1:KZ),:)+...
        [ones(KZ,5*N+2),zeros(KZ,N),ones(KZ,K+KZ+N*T)];
  % kappa
HPattern((6*N+2+K+KZ+1):end,:)=HPattern((6*N+2+K+KZ+1):end,:)+...
    [repmat([repmat(eye(N),1,2),ones(N,2*N+2),eye(N),zeros(N),ones(N,K+KZ)],T,1),eye(N*T)];
HPattern=1*(HPattern>0);

HPattern=sparse(HPattern);

clear n

% options
options=optimset('Display','off','JacobPattern',JPattern,'HessPattern',HPattern);


%% starting values

% "short" replication
load(['dR',num2str(100*diffRadius),'nl_replication.mat'],'MAP')
pars0=MAP;

% % "long" replication
% load(['dR',num2str(100*diffRadius),'nl_replication.mat'],'srng')
% rng(srng)
% alpha=normrnd(a0,om_a,N,1);
% theta=normrnd(th0,om_th,2,1);
% beta=[normrnd(b0A,om_b,N,1);normrnd(b0D,om_b,N,1)];
% v=gamrnd(s_v,1/d_v,N,1).^(-1);
% f=normrnd(f0,om_f,N,1);
% vs=gamrnd(s_vs,1/d_vs,N,1).^(-1);
% gamma=2*rand(KZ,1);
% xi=2*(-1).^(rand(K,1)<0.5).*rand(K,1);
% kppa=zeros(N*T,1);
% pars0=[alpha;theta;beta;v;f;vs;xi;gamma;kppa];
% clear alpha theta beta v f vs xi gamma kppa 


%% maximum-a-posteriori estimator
[MAP,lp,exit]=knitromatlab(@(x)log_posterior(x,D,M,X,Z,...
    a0,om_a,th0,om_th,b0A,b0D,om_b,s_v,d_v,f0,om_f,s_vs,d_vs),pars0,[],[],...
  [],[],[-Inf(3*N+2,1);zeros(N,1);-Inf(N,1);zeros(N,1);-Inf(K,1);zeros(KZ,1);-Inf(N*T,1)],...
  Inf(length(pars0),1),@(x)eq_con_jac_NOlearn(x,y,D,M,X,Z,sgi_nodes,sgi_weights,ep_sd,Sigma),[],...
  options,'knitro.opt');


%%
clearvars -except diffRadius MAP exit
save(['MAPdR',num2str(100*diffRadius),'nl.mat'])



